from torchvision import models
from torchvision import transforms
from PIL import Image
import torch


def main():
    dir_models()


def dir_models():
    dirs = dir(models)
    # print("dirs=", dirs)
    # alexNet()
    resnet = resNet100()

    preprocess = transfer()
    img = load_image()
    img_t = preprocess(img)
    batch_t = torch.unsqueeze(img_t, 0)
    resnet.eval()
    out = resnet(batch_t)
    print("", out)


def alexNet():
    alexnet = models.AlexNet()


def resNet100():
    resnet = models.resnet101(pretrained=True)
    # print("Resnet101=", resnet)
    return resnet


def transfer():
    preprocess = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        )])
    return preprocess


def load_image():
    img = Image.open("../data/p1ch2/bobby.jpg")
    # img.show()
    return img


if __name__ == '__main__':
    main()
